#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
健康检查脚本
功能：检查服务状态、系统资源，并在发现问题时触发恢复机制
作者：Ansible自动生成
"""

import os
import sys
import json
import time
import yaml
import logging
import requests
import subprocess
import psutil
from datetime import datetime, timedelta
from typing import Dict, List, Any, Optional


class HealthChecker:
    """健康检查器主类"""
    
    def __init__(self, config_file: str = '{{ failover_config_dir }}/config.yml'):
        self.config_file = config_file
        self.config = self.load_config()
        self.logger = self.setup_logging()
        self.state_manager = StateManager(self.config['global']['state_file'])
        self.alert_sender = AlertSender(self.config.get('alerts', {}), self.logger)
        
    def load_config(self) -> Dict[str, Any]:
        """加载配置文件"""
        try:
            with open(self.config_file, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        except Exception as e:
            print(f"Error loading config file {self.config_file}: {e}")
            sys.exit(1)
    
    def setup_logging(self) -> logging.Logger:
        """设置日志记录"""
        log_file = self.config['global']['log_file']
        log_level = getattr(logging, self.config['global']['log_level'].upper())
        
        # 确保日志目录存在
        os.makedirs(os.path.dirname(log_file), exist_ok=True)
        
        # 配置日志格式
        logging.basicConfig(
            level=log_level,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file),
                logging.StreamHandler(sys.stdout)
            ]
        )
        
        return logging.getLogger(__name__)
    
    def check_http_service(self, service: Dict[str, Any]) -> bool:
        """检查HTTP服务状态"""
        try:
            url = service['url']
            method = service.get('method', 'GET')
            timeout = service.get('timeout', 10)
            expected_status = service.get('expected_status', 200)
            expected_content = service.get('expected_content')
            
            self.logger.debug(f"Checking HTTP service: {url}")
            
            response = requests.request(method, url, timeout=timeout)
            
            # 检查状态码
            if response.status_code != expected_status:
                self.logger.warning(f"HTTP service {service['name']} returned status {response.status_code}, expected {expected_status}")
                return False
            
            # 检查响应内容（如果配置了）
            if expected_content and expected_content not in response.text:
                self.logger.warning(f"HTTP service {service['name']} response does not contain expected content")
                return False
            
            self.logger.debug(f"HTTP service {service['name']} is healthy")
            return True
            
        except requests.exceptions.RequestException as e:
            self.logger.error(f"HTTP service {service['name']} check failed: {e}")
            return False
        except Exception as e:
            self.logger.error(f"Unexpected error checking HTTP service {service['name']}: {e}")
            return False
    
    def check_tcp_service(self, service: Dict[str, Any]) -> bool:
        """检查TCP服务状态"""
        import socket
        
        try:
            host = service['host']
            port = service['port']
            timeout = service.get('timeout', 5)
            
            self.logger.debug(f"Checking TCP service: {host}:{port}")
            
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(timeout)
            result = sock.connect_ex((host, port))
            sock.close()
            
            if result == 0:
                self.logger.debug(f"TCP service {service['name']} is healthy")
                return True
            else:
                self.logger.warning(f"TCP service {service['name']} connection failed")
                return False
                
        except Exception as e:
            self.logger.error(f"TCP service {service['name']} check failed: {e}")
            return False
    
    def check_systemd_service(self, service: Dict[str, Any]) -> bool:
        """检查systemd服务状态"""
        try:
            service_name = service['service_name']
            
            self.logger.debug(f"Checking systemd service: {service_name}")
            
            result = subprocess.run(
                ['systemctl', 'is-active', service_name],
                capture_output=True,
                text=True
            )
            
            if result.returncode == 0 and result.stdout.strip() == 'active':
                self.logger.debug(f"Systemd service {service['name']} is healthy")
                return True
            else:
                self.logger.warning(f"Systemd service {service['name']} is not active")
                return False
                
        except Exception as e:
            self.logger.error(f"Systemd service {service['name']} check failed: {e}")
            return False
    
    def check_process(self, service: Dict[str, Any]) -> bool:
        """检查进程状态"""
        try:
            process_name = service['process_name']
            
            self.logger.debug(f"Checking process: {process_name}")
            
            for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
                try:
                    if process_name in proc.info['name'] or any(process_name in cmd for cmd in proc.info['cmdline']):
                        self.logger.debug(f"Process {service['name']} is running (PID: {proc.info['pid']})")
                        return True
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    continue
            
            self.logger.warning(f"Process {service['name']} is not running")
            return False
            
        except Exception as e:
            self.logger.error(f"Process {service['name']} check failed: {e}")
            return False
    
    def check_system_resources(self) -> Dict[str, bool]:
        """检查系统资源状态"""
        results = {}
        system_checks = self.config.get('system_checks', {})
        
        # CPU检查
        if system_checks.get('cpu', {}).get('enabled', True):
            cpu_threshold = system_checks['cpu'].get('threshold', 80)
            cpu_usage = psutil.cpu_percent(interval=1)
            results['cpu'] = cpu_usage < cpu_threshold
            if not results['cpu']:
                self.logger.warning(f"High CPU usage: {cpu_usage}% (threshold: {cpu_threshold}%)")
            else:
                self.logger.debug(f"CPU usage: {cpu_usage}%")
        
        # 内存检查
        if system_checks.get('memory', {}).get('enabled', True):
            memory_threshold = system_checks['memory'].get('threshold', 85)
            memory = psutil.virtual_memory()
            memory_usage = memory.percent
            results['memory'] = memory_usage < memory_threshold
            if not results['memory']:
                self.logger.warning(f"High memory usage: {memory_usage}% (threshold: {memory_threshold}%)")
            else:
                self.logger.debug(f"Memory usage: {memory_usage}%")
        
        # 磁盘检查
        if system_checks.get('disk', {}).get('enabled', True):
            disk_threshold = system_checks['disk'].get('threshold', 90)
            disk_paths = system_checks['disk'].get('paths', ['/'])
            results['disk'] = True
            
            for path in disk_paths:
                try:
                    disk = psutil.disk_usage(path)
                    disk_usage = disk.percent
                    if disk_usage >= disk_threshold:
                        results['disk'] = False
                        self.logger.warning(f"High disk usage on {path}: {disk_usage}% (threshold: {disk_threshold}%)")
                    else:
                        self.logger.debug(f"Disk usage on {path}: {disk_usage}%")
                except Exception as e:
                    self.logger.error(f"Error checking disk usage for {path}: {e}")
                    results['disk'] = False
        
        # 系统负载检查
        if system_checks.get('load', {}).get('enabled', True):
            load_threshold = system_checks['load'].get('threshold', 5.0)
            load_avg = psutil.getloadavg()[0]  # 1分钟平均负载
            results['load'] = load_avg < load_threshold
            if not results['load']:
                self.logger.warning(f"High system load: {load_avg} (threshold: {load_threshold})")
            else:
                self.logger.debug(f"System load: {load_avg}")
        
        return results
    
    def check_service(self, service: Dict[str, Any]) -> bool:
        """检查单个服务状态"""
        service_type = service.get('type', 'http')
        
        if service_type == 'http':
            return self.check_http_service(service)
        elif service_type == 'tcp':
            return self.check_tcp_service(service)
        elif service_type == 'systemd':
            return self.check_systemd_service(service)
        elif service_type == 'process':
            return self.check_process(service)
        else:
            self.logger.error(f"Unknown service type: {service_type}")
            return False
    
    def run_custom_check(self, check: Dict[str, Any]) -> bool:
        """运行自定义检查脚本"""
        try:
            script = check['script']
            timeout = check.get('timeout', 30)
            
            self.logger.debug(f"Running custom check: {check['name']}")
            
            result = subprocess.run(
                script,
                shell=True,
                capture_output=True,
                text=True,
                timeout=timeout
            )
            
            if result.returncode == 0:
                self.logger.debug(f"Custom check {check['name']} passed")
                return True
            else:
                self.logger.warning(f"Custom check {check['name']} failed: {result.stderr}")
                return False
                
        except subprocess.TimeoutExpired:
            self.logger.error(f"Custom check {check['name']} timed out")
            return False
        except Exception as e:
            self.logger.error(f"Custom check {check['name']} error: {e}")
            return False
    
    def attempt_recovery(self, service: Dict[str, Any]) -> bool:
        """尝试恢复服务"""
        recovery_config = service.get('recovery', {})
        
        if not recovery_config.get('enabled', True):
            self.logger.info(f"Recovery disabled for service {service['name']}")
            return False
        
        # 检查恢复冷却时间
        if not self.state_manager.can_attempt_recovery(service['name']):
            self.logger.info(f"Service {service['name']} is in recovery cooldown")
            return False
        
        max_attempts = recovery_config.get('max_attempts', 3)
        current_attempts = self.state_manager.get_recovery_attempts(service['name'])
        
        if current_attempts >= max_attempts:
            self.logger.error(f"Service {service['name']} has exceeded maximum recovery attempts ({max_attempts})")
            return False
        
        try:
            restart_command = recovery_config.get('restart_command', f'systemctl restart {service["name"]}')
            
            self.logger.info(f"Attempting to recover service {service['name']} (attempt {current_attempts + 1}/{max_attempts})")
            
            # 执行恢复命令
            result = subprocess.run(
                restart_command,
                shell=True,
                capture_output=True,
                text=True,
                timeout=60
            )
            
            if result.returncode == 0:
                self.logger.info(f"Recovery command executed successfully for {service['name']}")
                self.state_manager.record_recovery_attempt(service['name'], True)
                
                # 等待一段时间后重新检查
                time.sleep(10)
                if self.check_service(service):
                    self.logger.info(f"Service {service['name']} recovered successfully")
                    self.state_manager.reset_recovery_attempts(service['name'])
                    return True
                else:
                    self.logger.warning(f"Service {service['name']} still unhealthy after recovery attempt")
                    return False
            else:
                self.logger.error(f"Recovery command failed for {service['name']}: {result.stderr}")
                self.state_manager.record_recovery_attempt(service['name'], False)
                return False
                
        except subprocess.TimeoutExpired:
            self.logger.error(f"Recovery command timed out for {service['name']}")
            self.state_manager.record_recovery_attempt(service['name'], False)
            return False
        except Exception as e:
            self.logger.error(f"Recovery attempt failed for {service['name']}: {e}")
            self.state_manager.record_recovery_attempt(service['name'], False)
            return False
    
    def run_health_checks(self):
        """运行所有健康检查"""
        self.logger.info("Starting health checks")
        
        # 检查服务
        services = self.config.get('services', [])
        for service in services:
            service_name = service['name']
            
            try:
                is_healthy = self.check_service(service)
                
                if is_healthy:
                    self.logger.info(f"Service {service_name} is healthy")
                    self.state_manager.update_service_state(service_name, 'healthy')
                else:
                    self.logger.warning(f"Service {service_name} is unhealthy")
                    self.state_manager.update_service_state(service_name, 'unhealthy')
                    
                    # 发送告警
                    self.alert_sender.send_alert(
                        f"Service {service_name} is unhealthy",
                        f"Health check failed for service {service_name}",
                        'warning'
                    )
                    
                    # 尝试恢复
                    if self.config.get('recovery', {}).get('enabled', True):
                        recovery_success = self.attempt_recovery(service)
                        if recovery_success:
                            self.alert_sender.send_alert(
                                f"Service {service_name} recovered",
                                f"Service {service_name} has been successfully recovered",
                                'info'
                            )
                        else:
                            self.alert_sender.send_alert(
                                f"Service {service_name} recovery failed",
                                f"Failed to recover service {service_name}",
                                'critical'
                            )
                            
            except Exception as e:
                self.logger.error(f"Error checking service {service_name}: {e}")
                self.state_manager.update_service_state(service_name, 'error')
        
        # 检查系统资源
        system_results = self.check_system_resources()
        if system_results is not None:
            for resource, is_healthy in system_results.items():
                if not is_healthy:
                    self.alert_sender.send_alert(
                        f"High {resource} usage detected",
                        f"System {resource} usage is above threshold",
                        'warning'
                    )
        
        # 运行自定义检查
        custom_checks = self.config.get('custom_checks', []) or []
        for check in custom_checks:
            if check.get('enabled', True):
                try:
                    is_healthy = self.run_custom_check(check)
                    if not is_healthy:
                        self.alert_sender.send_alert(
                            f"Custom check {check['name']} failed",
                            f"Custom health check {check['name']} returned failure",
                            'warning'
                        )
                except Exception as e:
                    self.logger.error(f"Error running custom check {check['name']}: {e}")
        
        self.logger.info("Health checks completed")


class StateManager:
    """状态管理器"""
    
    def __init__(self, state_file: str):
        self.state_file = state_file
        self.state = self.load_state()
    
    def load_state(self) -> Dict[str, Any]:
        """加载状态文件"""
        try:
            if os.path.exists(self.state_file):
                with open(self.state_file, 'r') as f:
                    return json.load(f)
        except Exception:
            pass
        
        return {
            'services': {},
            'recovery_attempts': {},
            'last_recovery': {},
            'created_at': datetime.now().isoformat()
        }
    
    def save_state(self):
        """保存状态文件"""
        try:
            os.makedirs(os.path.dirname(self.state_file), exist_ok=True)
            with open(self.state_file, 'w') as f:
                json.dump(self.state, f, indent=2)
        except Exception as e:
            print(f"Error saving state file: {e}")
    
    def update_service_state(self, service_name: str, state: str):
        """更新服务状态"""
        self.state['services'][service_name] = {
            'state': state,
            'last_check': datetime.now().isoformat()
        }
        self.save_state()
    
    def get_recovery_attempts(self, service_name: str) -> int:
        """获取恢复尝试次数"""
        return self.state['recovery_attempts'].get(service_name, 0)
    
    def record_recovery_attempt(self, service_name: str, success: bool):
        """记录恢复尝试"""
        if service_name not in self.state['recovery_attempts']:
            self.state['recovery_attempts'][service_name] = 0
        
        if not success:
            self.state['recovery_attempts'][service_name] += 1
        
        self.state['last_recovery'][service_name] = {
            'timestamp': datetime.now().isoformat(),
            'success': success
        }
        
        self.save_state()
    
    def reset_recovery_attempts(self, service_name: str):
        """重置恢复尝试次数"""
        self.state['recovery_attempts'][service_name] = 0
        self.save_state()
    
    def can_attempt_recovery(self, service_name: str) -> bool:
        """检查是否可以尝试恢复（冷却时间）"""
        last_recovery = self.state['last_recovery'].get(service_name)
        if not last_recovery:
            return True
        
        last_time = datetime.fromisoformat(last_recovery['timestamp'])
        cooldown = timedelta(seconds=300)  # 5分钟冷却时间
        
        return datetime.now() - last_time > cooldown


class AlertSender:
    """告警发送器"""
    
    def __init__(self, alert_config: Dict[str, Any], logger: logging.Logger):
        self.config = alert_config
        self.logger = logger
    
    def send_alert(self, title: str, message: str, level: str = 'info'):
        """发送告警"""
        if not self.config.get('enabled', True):
            return
        
        alert_data = {
            'timestamp': datetime.now().isoformat(),
            'title': title,
            'message': message,
            'level': level,
            'hostname': os.uname().nodename
        }
        
        # 记录到日志文件
        if self.config.get('log', {}).get('enabled', True):
            self._log_alert(alert_data)
        
        # 发送Webhook告警
        if self.config.get('webhook', {}).get('enabled', False):
            self._send_webhook_alert(alert_data)
        
        # 发送邮件告警
        if self.config.get('email', {}).get('enabled', False):
            self._send_email_alert(alert_data)
    
    def _log_alert(self, alert_data: Dict[str, Any]):
        """记录告警到日志文件"""
        try:
            log_file = self.config['log'].get('file', '/var/log/failover/alerts.log')
            os.makedirs(os.path.dirname(log_file), exist_ok=True)
            
            with open(log_file, 'a') as f:
                f.write(json.dumps(alert_data) + '\n')
        except Exception as e:
            self.logger.error(f"Failed to log alert: {e}")
    
    def _send_webhook_alert(self, alert_data: Dict[str, Any]):
        """发送Webhook告警"""
        try:
            webhook_url = self.config['webhook']['url']
            timeout = self.config['webhook'].get('timeout', 10)
            
            response = requests.post(
                webhook_url,
                json=alert_data,
                timeout=timeout
            )
            
            if response.status_code == 200:
                self.logger.debug("Webhook alert sent successfully")
            else:
                self.logger.warning(f"Webhook alert failed with status {response.status_code}")
                
        except Exception as e:
            self.logger.error(f"Failed to send webhook alert: {e}")
    
    def _send_email_alert(self, alert_data: Dict[str, Any]):
        """发送邮件告警"""
        try:
            import smtplib
            from email.mime.text import MIMEText
            from email.mime.multipart import MIMEMultipart
            
            email_config = self.config['email']
            
            msg = MIMEMultipart()
            msg['From'] = email_config['from_email']
            msg['To'] = ', '.join(email_config['to_emails'])
            msg['Subject'] = f"[{alert_data['level'].upper()}] {alert_data['title']}"
            
            body = f"""
            告警时间: {alert_data['timestamp']}
            主机名: {alert_data['hostname']}
            告警级别: {alert_data['level']}
            告警标题: {alert_data['title']}
            告警内容: {alert_data['message']}
            """
            
            msg.attach(MIMEText(body, 'plain'))
            
            server = smtplib.SMTP(email_config['smtp_server'], email_config['smtp_port'])
            server.starttls()
            server.login(email_config['username'], email_config['password'])
            
            text = msg.as_string()
            server.sendmail(email_config['from_email'], email_config['to_emails'], text)
            server.quit()
            
            self.logger.debug("Email alert sent successfully")
            
        except Exception as e:
            self.logger.error(f"Failed to send email alert: {e}")


def main():
    """主函数"""
    try:
        checker = HealthChecker()
        checker.run_health_checks()
    except KeyboardInterrupt:
        print("\nHealth check interrupted by user")
        sys.exit(0)
    except Exception as e:
        import traceback
        print(f"Health check failed: {e}")
        print(f"Traceback: {traceback.format_exc()}")
        sys.exit(1)


if __name__ == '__main__':
    main()